import random
import torch
from torch.utils.data import Dataset, DataLoader, Subset, random_split

torch.manual_seed(42)

class DS(Dataset):
    def __init__(self):
        super().__init__()
        self.x = torch.arange(105).reshape(-1, 1).float()
        self.y = 0.6 * self.x + 2       # y = wx+b
        self.y += torch.normal(0, 1, self.y.shape)

    def __len__(self):
        return len(self.y)

    def __getitem__(self, idx):
        return self.x[idx], self.y[idx]

ds = DS()
batch_size = 16
dl = DataLoader(ds, batch_size=batch_size, shuffle=True)
print(f"取到数据集下标为0的x和y {ds[0]}")
print(f"随机批量加载的dl的总批次是 {len(dl)}")
print(f"取到第一批的数据是 {next(enumerate(dl))}")

# Subset 数据集的子集类
# 用于从数据集中选中指定下标索引的子集，这对于分割数据集比较有用
s1 = Subset(ds, [0, 2, 3])
print(s1[0])    # (tensor([0.]), tensor([3.9269]))
print(s1[1])    # (tensor([2.]), tensor([4.1007]))
print(s1[2])    # (tensor([3.]), tensor([1.6945]))
# 随机采样
idx = torch.arange(len(ds)).tolist()
random_idx = random.sample(idx, k=round(len(ds)*0.2))   # 随机采样20%
s2 = Subset(ds, random_idx)
print(len(s2))  # 21
print(s2[0])    # (tensor([61.]), tensor([40.4446]))
print(s2[1])    # (tensor([95.]), tensor([58.6964]))

# torch.utils.data.random_split是Dataset的一个子类，用于随机划分数据集
train_ds_rate = 0.8
val_ds_rate = 0.1
test_ds_rate = 0.1
train_ds_len = round(len(ds)*train_ds_rate)
val_ds_len = round(len(ds)*val_ds_rate)
test_ds_len = len(ds) - train_ds_len - val_ds_len
train_ds, val_ds, test_ds = random_split(ds, [train_ds_len, val_ds_len, test_ds_len])
print(len(train_ds))
print(len(val_ds))
print(len(test_ds))

